package com.vbobot.demo.spring.cloud.gateway.ribbon.rules;

import brave.propagation.ExtraFieldPropagation;
import com.netflix.loadbalancer.Server;
import com.netflix.loadbalancer.ZoneAvoidanceRule;
import com.vbobot.demo.spring.cloud.gateway.ribbon.rules.hashing.ConsistentHashLoadBalancer;
import com.vbobot.demo.spring.cloud.gateway.ribbon.rules.hashing.MurmurHashStrategyImpl;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import org.apache.commons.lang3.StringUtils;

/**
 * https://github.com/alejandro-du/vaadin-microservices-demo/blob/263cfd4da791c038042c4efbe0cab40ce831273c/proxy-server/src/main/java/com/example/StickySessionRule.java#L41
 *
 * @author Bobo
 * @date 2022/11/1
 */
public class StickySessionRule extends ZoneAvoidanceRule {

    public static final String X_COMPANY_ID = "x-company-id";
    private final ConsistentHashLoadBalancer consistentHashLoadBalancer = new ConsistentHashLoadBalancer(
            new MurmurHashStrategyImpl());

    @Override
    public Server choose(Object key) {
        final List<Server> reachableServers = getLoadBalancer().getReachableServers();
        if (reachableServers == null || reachableServers.size() <= 1) {
            return superChoose(key);
        }

        final List<Server> availableServers = reachableServers.stream()
                .filter(item -> item.isAlive() && item.isReadyToServe())
                .collect(Collectors.toList());
        if (availableServers.size() <= 1) {
            return superChoose(key);
        }

        final Server chooseServer = chooseServerIndex(availableServers);
        if (chooseServer != null) {
            return chooseServer;
        }

        return superChoose(key);
    }

    private Server superChoose(Object key) {
        return super.choose(key);
    }

    private Server chooseServerIndex(List<Server> availableServers) {
        final String companyId = ExtraFieldPropagation.get(X_COMPANY_ID);
        if (StringUtils.isBlank(companyId)) {
            return null;
        }

        final List<InnerServer> serverList = availableServers.stream()
                .map(InnerServer::new)
                .collect(Collectors.toList());

        return Optional.ofNullable(consistentHashLoadBalancer.select(serverList, () -> companyId))
                .map(item -> item.server)
                .orElse(null);
    }

    private static class InnerServer implements com.vbobot.demo.spring.cloud.gateway.ribbon.rules.hashing.Server {

        private final Server server;

        public InnerServer(Server server) {
            this.server = server;
        }

        @Override
        public String getUrl() {
            return server.getMetaInfo().getInstanceId();
        }
    }
}
